import pickle 
import json
import numpy as np
from pathlib import Path
import os 
import sys
from utils import load_json,gen_dict_from_raw_json,Sensor_data

data_root = "/home/yyj/dl_dataset/nuscenes_mini/"
json_file_dir = os.path.join(data_root,"v1.0-mini")

if not os.path.exists(json_file_dir):
    raise Exception("{} dir not exist.".format(json_file_dir))

ego_pose_file_path            = os.path.join(json_file_dir,"ego_pose.json")
sample_file_path              = os.path.join(json_file_dir,"sample.json")
sample_data_file_path         = os.path.join(json_file_dir,"sample_data.json")
scene_file_path               = os.path.join(json_file_dir,"scene.json")
calibrated_sensor_file_path   = os.path.join(json_file_dir,"calibrated_sensor.json")
sensor_file_path              = os.path.join(json_file_dir,"sensor.json")

ego_pose_json          = load_json(ego_pose_file_path)
sample_json            = load_json(sample_file_path)
scene_json             = load_json(scene_file_path)
sample_data_json       = load_json(sample_data_file_path)
calibrated_sensor_json = load_json(calibrated_sensor_file_path)
sensor_json            = load_json(sensor_file_path)

#需要将每个 json 文件中的 list  改成 dict 的形式进行索引
# print("sample_data json : ")
# print(sample_data_json[1])  
# print('sample json:')
# print(sample_json[1])
# print('calibrated_sensor json:')
# print(calibrated_sensor_json[1])
# print('scene json: ')
# print(scene_json[1])
#! 需要使用 sample_data json 中的内容 ， 来读取非关键帧的消息
"""
    0. 将 scene json 中的内容生成dict
    1. 将 sample json 中生成 token -> sample dict 
    2. 将sample data中的内容 根据sample 中的内容找到对应的scene token 从而获得 name 这个属性
    3. 是用 dict {scene_name : [sample_data]} 这个数据结构来存储当前 scene 中所有的sample data
"""

scene_token_data_dict             = gen_dict_from_raw_json(scene_json)
sample_token_data_dcit            = gen_dict_from_raw_json(sample_json)
ego_pose_token_data_dict          = gen_dict_from_raw_json(ego_pose_json)
calibrated_sensor_token_data_dict = gen_dict_from_raw_json(calibrated_sensor_json)
sensor_token_data_dcit            = gen_dict_from_raw_json(sensor_json)
scene_sample_data_dict            = {scene_token : [] for scene_token in scene_token_data_dict.keys()}

for ind,sample_data in enumerate(sample_data_json):
    sample_token = sample_data['sample_token']
    scene_token = sample_token_data_dcit[sample_token]['scene_token']
    stamp = sample_data['timestamp'] / 1e6
    ref_filename = sample_data['filename']
    ego_pose = ego_pose_token_data_dict[sample_data['ego_pose_token']]
    calibrated_sensor = calibrated_sensor_token_data_dict[sample_data['calibrated_sensor_token']]
    current_sensor_name = sensor_token_data_dcit[calibrated_sensor['sensor_token']]['channel']

    current_sensor_data = Sensor_data(current_sensor_name,ref_filename,stamp,ego_pose,sample_data['token'])
    scene_sample_data_dict[scene_token].append(current_sensor_data)


for scene_token in scene_sample_data_dict.keys():
    scene_sample_data_dict[scene_token] = sorted(scene_sample_data_dict[scene_token],key=lambda x:x.stamp)

    # for sample_data in scene_sample_datas:
    #     print("stmap : {}    sensor_name : {}  ref_file_path = {}".format(sample_data.stamp,sample_data.sensor_name,sample_data.ref_file_path))
        # print(sample_data.ego_pose.rot ,sample_data.ego_pose.t )

with open('scene_sample_data_dcit.pickle', 'wb') as f:
    pickle.dump(scene_sample_data_dict, f)

print('scene_sample_data_dict saved.')